import os

from dotenv import load_dotenv
from openai import OpenAI
import openpyxl

# 加载.env文件（OpenAI key）
load_dotenv()
BASE_URL = os.getenv("BASE_URL")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# 创建OpenAI连接
client = OpenAI(
    base_url=BASE_URL,
    api_key=OPENAI_API_KEY
)


def write_txt(text):
    """
    内容写入到txt中
    :param text:
    :return:
    """
    file_path = f'./local_data/output.txt'
    with open(file_path, 'a', encoding='utf8') as file:
        file.write(text)


# 读取excel
workbook = openpyxl.load_workbook('./local_data/招标采购标的物信息提取训练数据.xlsx')
worksheet = workbook.worksheets[1]

row = 1748
while row < worksheet.max_row + 1:
    project_name = worksheet.cell(row=row, column=9).value
    prompt = f"对如下句子进行文本分类，选项有工程、服务、采购：{project_name}。只回答是哪个类别，不需要解释。"
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo-0125",
            messages=[
                {"role": "user",
                 "content": prompt},
            ],
            temperature=0,
        )
        project_class = response.choices[0].message.content
        project_class.replace('。', '')
        worksheet.cell(row=row, column=10).value = project_class
        write_txt(f'{row}-{project_name}-{project_class}\n')
        print(f'{row}-{project_name}-{project_class}')
    except Exception as e:
        print(e)
        row -= 1
    row += 1

workbook.close()
